# -*- coding: utf-8 -*-
"""
Created on Wed Apr 13 10:49:51 2022

@author: 04566
"""
import paddle
import paddle.nn as pdnn
from paddle.vision.models import mobilenet_v1

class mobilenet_encoder(pdnn.Layer):
    def __init__(self,pretrained=True):
        super(mobilenet_encoder,self).__init__()
        self.encoder=mobilenet_v1(pretrained)
        self.seq_list=[]
        chs=[(128,64),(256,128),(512,256),(1024,512)]
        for ch in chs:
            self.seq=pdnn.Sequential(pdnn.Conv2D(in_channels=ch[0],out_channels=ch[1],kernel_size=1),\
                                     pdnn.Upsample(scale_factor=2),\
                                     pdnn.Conv2D(in_channels=ch[1],out_channels=ch[1],kernel_size=3,padding=1,padding_mode='reflect'),
                                     pdnn.ELU(),\
                                     pdnn.MaxPool2D(2),
                                     pdnn.Conv2D(in_channels=ch[1],out_channels=ch[1],kernel_size=1))
            self.seq_list.append(self.seq)
        self.seq_list=pdnn.LayerList(self.seq_list)
        '''
        self.conv1=pdnn.Conv2D(in_channels=128,out_channels=64,kernel_size=1)
        self.conv2=pdnn.Conv2D(in_channels=256,out_channels=128,kernel_size=1)
        self.conv3=pdnn.Conv2D(in_channels=512,out_channels=256,kernel_size=1)
        self.conv4=pdnn.Conv2D(in_channels=1024,out_channels=512,kernel_size=1)
        '''
    def forward(self,x):
        feature=[]
        x=self.encoder.conv1(x)
        i=1
        for dws in self.encoder.dwsl:
            x = dws(x)
            if i==1:
                feature.append(x)
            elif i==3:
                feature.append(self.seq_list[0](x))
            elif i==5:
                feature.append(self.seq_list[1](x))
            elif i==7:
                feature.append(self.seq_list[2](x))
            i+=1
        feature.append(self.seq_list[3](x))
        
        return feature
'''    
x=paddle.ones((1,3,192,640))
model=mobilenet_encoder()
y=model(x)
print('0',y[0].shape)
print('1',y[1].shape)
print('2',y[2].shape)
print('3',y[3].shape)
print('4',y[4].shape)
0 [1, 64, 96, 320]
1 [1, 128, 48, 160]
2 [1, 256, 24, 80]
3 [1, 512, 12, 40]
4 [1, 1024, 6, 20]
'''